import argparse
import torch
from torchvision import transforms

from retinanet import model
from retinanet.dataloader import CSVDataset, Resizer, Normalizer
from retinanet import csv_eval
import numpy as np
assert torch.__version__.split('.')[0] == '1'

print('CUDA available: {}'.format(torch.cuda.is_available()))

iou_thresh = 0.3
score_thresh = 0.05

def main(args=None):
    parser = argparse.ArgumentParser(description='Simple training script for training a RetinaNet network.')

    parser.add_argument('--dataset', help='Dataset type, must be one of csv or coco.', default='csv')
    parser.add_argument('--coco_path', help='Path to COCO directory')
    parser.add_argument('--csv_val', help='Path to file containing validation annotations (optional, see readme)',
                        default='../data/csv_label_new/zhengwei/val_zhengwei_box.csv')
    # parser.add_argument('--model_path', help='Path to model', type=str,
    #                     default='../model_save/detection/cewei/res18/_retinanet_70.pth')
    parser.add_argument('--model_path', help='Path to model', type=str,
                        # default='../model_save/detection/cewei/eff_b0/_retinanet_138.pth')
                        default='../model_save/new/detection/zhengwei/new_efficientdet_b3/_retinanet_68.pth')

    parser.add_argument('--csv_classes', help='Path to file containing class list (see readme)',
                        default='../data/csv_label/class.csv')
    parser.add_argument('--depth', help='Resnet depth, must be one of 18, 34, 50, 101, 152', default='new_efficientdet_b3')


    parser = parser.parse_args(args)

    dataset_val = CSVDataset(train_file=parser.csv_val, class_list=parser.csv_classes,
                                     transform=transforms.Compose([Normalizer(), Resizer()]))

    # Create the model
    if parser.depth == 18:
        retinanet = model.resnet18(num_classes=1, pretrained=True)
    elif parser.depth == 34:
        retinanet = model.resnet34(num_classes=1, pretrained=True)
    elif parser.depth == 50:
        retinanet = model.resnet50(num_classes=1, pretrained=True)
    elif parser.depth == 101:
        retinanet = model.resnet101(num_classes=1, pretrained=True)
    elif parser.depth == 152:
        retinanet = model.resnet152(num_classes=1, pretrained=True)
    elif parser.depth == 'efficientnet-b0':
        retinanet = model.resnet18(pretrained=True, num_classes=1,
                                   effici_name='efficientnet-b0')
    elif parser.depth == 'new_efficientdet_b3':
        retinanet = model.new_efficientdet_b3(pretrained=True, num_classes=1, thresh=iou_thresh)
    else:
        raise ValueError('Unsupported model depth, must be one of 18, 34, 50, 101, 152')
    use_gpu = torch.cuda.is_available()

    if use_gpu:
        retinanet = retinanet.cuda()

    # load trained model weights
    trained_model_state_dict = torch.load(parser.model_path)
    trained_model_state_dict = {k.replace('module.', ''): v for k, v in trained_model_state_dict.items()}
    retinanet.load_state_dict(trained_model_state_dict)


    retinanet = torch.nn.DataParallel(retinanet).cuda()

    retinanet.training = False
    retinanet.eval()
    retinanet.module.freeze_bn()

    AP = csv_eval.evaluate(dataset_val, retinanet,iou_threshold=0.75,score_threshold=score_thresh,max_detections=100)
    mAP = np.array([AP[i][0] for i in range(len(AP))])
    print('mAP',mAP.mean())


if __name__ == '__main__':
    main()
